import contextlib
import os
import struct
import wave
import librosa
import numpy as np

# files 文件路径
# time 为语音总长度
def red_time(files, times):
    for path, dir_name, file_xm in os.walk(files):
        for name in file_xm:
            if name.endswith('wav'):
                path_wav = os.path.join(path, name)
                # with contextlib.closing(wave.open(path_wav, 'rb')) as f:
                #     frames = f.getnframes()
                #     rate = f.getframerate()
                #     duration = frames / float(rate)

                duration = librosa.get_duration(filename=path_wav)    
                # 判断当时语音长度
                if times > round(duration, 1):
                    durations = round(times - round(duration, 1), 1)
                    print(name)
                    print( durations)
                    # 调用添加空白语音
                    creata_wav(durations, path_wav, name)

def creata_wav(durations=None, path_wav=None, name=None):
    framerate = 44100
    sample_width = 2
    duration = durations
    frequency = 2000
    volume = 1000
    x = np.linspace(0, duration, num=int(duration * framerate))
    y = np.sin(0 * np.pi * frequency * x) * volume
    sine_wave = y
    sine = name.split('.')[0] + '空白.wav'
    with wave.open(sine, 'wb') as wf:
        wf.setnchannels(1)
        wf.setframerate(framerate)
        wf.setsampwidth(sample_width)
        for i in sine_wave:
            data = struct.pack('<h', int(i))
            wf.writeframesraw(data)
        wf.close()
    # 调用写功能
    red_wav(path_wav, sine)

def red_wav(path_wav, sine):
    with open(sine, 'rb') as file:
        data1 = file.read()
    with open(path_wav, 'rb') as file:
        data2 = file.read()
    data_info = data1[:44]  # 复制帧头参考
    data_out = data1[44:] + data2[44:]  # 将两个音频的数据帧合并（都是相同格式）
    data_info = data_info[:4] + struct.pack('<L', len(data_out) + 44) + data_info[8:]  # 更新WAV文件的总byte数（两个文件数据帧和+44）
    data_info = data_info[:40] + struct.pack('<L', len(data_out)) + data_info[44:]  # 更新WAV文件的数据byte数（两个文件数据帧和）
    # *** 生成合并后的WAV文件 *** #
    with open(path_wav, 'wb') as f:
        f.write(data_info + data_out)
    print('完成：', path_wav)

if __name__ == "__main__":
    file_path = r'C:\Users\zz\Desktop\热词唤醒\keyword-spotting_2\data_CTC\words\shanghai'
    file_time = 1
    red_time(file_path, file_time)
    print('完成')
